from smolagents import ToolCallingAgent, LiteLLMModel, tool, GradioUI
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from VectorDBConnector import get_completion,get_embeddings,MyVectorDBConnector #addDocuments作为数据导入 

import logging
import os


# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

logger.info("正在初始化嵌入模型...")
from openai import OpenAI
client = OpenAI(
    api_key="key",
    base_url="http://localhost:11434/v1"
)
# 加载本地存储的向量数据库
logger.info("正在加载向量数据库...")
# 创建一个向量数据库对象
vectordb = MyVectorDBConnector(client,"corcName", get_embeddings)
logger.info("向量数据库加载成功")

model = LiteLLMModel(
    model_id="ollama/qwen3:32b",
    api_base="http://localhost:11434",
    api_key="your_ollama_api_key" #os.getenv("OPENROUTER_API_KEY"),  # 从环境变量中获取 API 密钥
)

@tool   
def retriever(query:str)->str:
    """
    根据用户的查询，执行向量数据库的相似性搜索，并返回结果的字符串表示形式。

    Args:
        query: 要查询的字符串。此字符串将用于在向量数据库中进行相似性搜索。

    """
    logger.info(f"正在查询: {query}")
    results = vectordb.search(query, k=5)  # k 是返回的结果数量
    logger.info("查询完成",results)

    # 将结果组合成一个字符串
    combined_results = "\n\n".join([f"资料{i+1}: {result.page_content}" for i, result in enumerate(results)])
    return combined_results


question = "你好"

rag_agent_prompt = f"""
根据你的知识库，回答以下问题。
请只回答问题，回答应该简洁且与问题相关。
如果你无法找到信息，不要放弃，尝试使用不同的参数再次调用你的 retriever 工具。
确保通过多次使用语义不同的查询来完全覆盖问题。
你的查询不应是问题，而是肯定形式的句子：例如，与其问"如何从 Hub 加载 bf16 模型？"，不如问"从 Hub 加载 bf16 权重"。

Question:
{question}"""

agent = ToolCallingAgent(tools=[retriever], model=model, add_base_tools=False)

agent.run(
    rag_agent_prompt,
)
